import uuid
from typing import List

from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import TextLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain.retrievers import MultiVectorRetriever
from langchain.schema import StrOutputParser, Document
from langchain.storage import InMemoryStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores.chroma import Chroma


def multi_vector_test():
    """
    每个文档存储多个向量通常是有益的。在许多用例中，这是有益的。 LangChain 有一个基础 MultiVectorRetriever，可以轻松查询此类设置。很多复杂性在于如何为每个文档创建多个向量。本笔记本涵盖了创建这些向量和使用 MultiVectorRetriever 的一些常见方法。

    为每个文档创建多个向量的方法包括：

    较小的块：将文档分割成较小的块，然后嵌入这些块（这是 ParentDocumentRetriever）。

    摘要：为每个文档创建摘要，将其与文档一起嵌入（或代替文档）。

    假设性问题：创建每个文档都适合回答的假设性问题，将这些问题与文档一起嵌入（或代替文档）。

    请注意，这还启用了另一种添加嵌入的方法 - 手动。这很棒，因为您可以显式添加导致文档恢复的问题或查询，从而为您提供更多控制
    """
    loaders = [
        TextLoader("files/paul_graham_essay.txt"),
        TextLoader("files/state_of_the_union.txt"),
    ]
    docs = []
    for l in loaders:
        docs.extend(l.load())
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000)
    docs = text_splitter.split_documents(docs)
    print(docs)
    return docs


def smaller_chunks():
    """
    通常，检索较大的信息块但嵌入较小的信息块可能很有用。这允许嵌入尽可能接近地捕获语义，但可以将尽可能多的上下文传递到下游。
    请注意，这就是 ParentDocumentRetriever 的作用。在这里，我们展示了幕后发生的事情
    """
    vectorstore = Chroma(collection_name="full_documents", embedding_function=OpenAIEmbeddings())
    store = InMemoryStore()
    id_key = "doc_id"
    retriever = MultiVectorRetriever(
        vectorstore=vectorstore,
        docsstore=store,
        id_key=id_key
    )
    docs = multi_vector_test()
    doc_ids = [str(uuid.uuid4()) for _ in docs]
    child_text_splitter = RecursiveCharacterTextSplitter(chunk_size=400)
    sub_docs = []
    for i, doc in enumerate(docs):
        _id = doc_ids[i]
        _sub_docs = child_text_splitter.split_documents([doc])
        for _doc in _sub_docs:
            _doc.metadata[id_key] = _id
        sub_docs.extend(_sub_docs)
    retriever.vectorstore.add_documents(sub_docs)
    retriever.docstore.mset(list(zip(doc_ids, docs)))
    # Vectorstore alone retrieves the small chunks
    print(retriever.vectorstore.similarity_search("justice breyer")[0])
    print(len(retriever.get_relevant_documents("justice breyer")[0].page_content))


def summary(docs: List):
    """
    通常，摘要可能能够更准确地提炼出某个块的内容，从而实现更好的检索。在这里，我们展示如何创建摘要，然后嵌入它们
    :return:
    """
    chain = (
            {"doc": lambda x: x.page_content}
            | ChatPromptTemplate.from_template("Summarize the following document:\n\n{doc}")
            | ChatOpenAI(max_retries=0)
            | StrOutputParser()
    )
    summaries = chain.batch(docs, {"max_concurrency": 5})
    # The vectorstore to use to index the child chunks
    vectorstore = Chroma(collection_name="summaries", embedding_function=OpenAIEmbeddings())
    # The storage layer for the parent documents
    store = InMemoryStore()
    id_key = "doc_id"
    # The retriever (empty to start)
    retriever = MultiVectorRetriever(
        vectorstore=vectorstore,
        docstore=store,
        id_key=id_key,
    )
    doc_ids = [str(uuid.uuid4()) for _ in docs]
    summary_docs = [
        Document(page_content=s, metadata={id_key: doc_ids[i]})
        for i, s in enumerate(summaries)
    ]
    retriever.vectorstore.add_documents(summary_docs)
    retriever.docstore.mset(list(zip(doc_ids, docs)))
    sub_docs = vectorstore.similarity_search("justice breyer")
    print(sub_docs[0])
    retrieved_docs = retriever.get_relevant_documents("justice breyer")
    print(len(retrieved_docs[0].page_content))

if __name__ == '__main__':
    smaller_chunks()
